import time
import logging
import requests
from typing import Optional, List, Dict, Mapping, Any

import langchain
from langchain.llms.base import LLM
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.cache import InMemoryCache
import jwt

logging.basicConfig(level=logging.INFO)
# 启动llm的缓存
langchain.llm_cache = InMemoryCache()



class ChatGLM_LLM(LLM):
    
    # 模型服务url
    url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
    model_name = "glm-4"

    def generate_token(apikey: str, exp_seconds: int):
        try:
            id, secret = apikey.split(".")
        except Exception as e:
            raise Exception("invalid apikey", e)

        payload = {
            "api_key": id,
            "exp": int(round(time.time() * 1000)) + exp_seconds * 1000,
            "timestamp": int(round(time.time() * 1000)),
        }

        return jwt.encode(
            payload,
            secret,
            algorithm="HS256",
            headers={"alg": "HS256", "sign_type": "SIGN"},
        )
        
    headers = {'Authorization': generate_token('bdc59e310deb29d48e6be230d487c518.n88YR9GP76XUePoL', 10000)}

    @property
    def _llm_type(self) -> str:
        return "chatglm"

    def _construct_query(self, prompt: str) -> Dict:
        query = {
            "messages":[
                {"role": "user",
                 "content": prompt}
            ],
            "model":self.model_name
        }
        return query

    def _call(self, prompt: str, stop: Optional[List[str]] = None,run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,) -> Any:
        query = self._construct_query(prompt=prompt)
        with requests.session() as session:
            resp = session.post(
                self.url,
                json=query,
                headers=self.headers
            )
            print(f"query:{query}")
            if resp.status_code == 200:
                resp_json = resp.json()
                choices  = resp_json.get("choices")
                if choices:
                    first_choice = choices[0]['message']['content']
                    return first_choice
                else:
                    print("Key 'choices' not found in response.")
            else:
                print(f"Request failed with status code {resp.status_code}: {resp.text}")
        return "请求失败"
        
